from climate_indices import indices
import os
import numpy as np
import math
import datetime

all_indices_root = os.path.abspath('result/all_indices')  # 临时站点数据存放
spei_root = os.path.abspath('result/multi_spei')
if not os.path.exists(spei_root):
    os.mkdir(spei_root)
multi_spei = os.path.abspath('result/multi_spei/multi_spei-12')# 复合SPEI值存放
if not os.path.exists(multi_spei):
    os.mkdir(multi_spei)
list = os.listdir(all_indices_root)

for i in range(0, len(list)):
# for i in range(0, 1):

    path = os.path.join(all_indices_root, list[i])
    if os.path.isfile(path):
        f = open(path, 'r')
        alllines = f.readlines()
        fp = os.path.join(multi_spei, 'Multi_SPEI-12_' + list[i][14:19] + '.txt')
        fw = open(fp, 'a+')
        pre_current = []  # 降水
        tem_current = []  # 温度
        for j in range(0, len(alllines)):
            try:
                alllines[j] = alllines[j].strip('\n').split(',')
            except:
                continue
            pre_current.append(float(alllines[j][3]))
            tem_current.append(float(alllines[j][5]))
        precips_np = np.array(pre_current).astype(np.float)
        tem_np = np.array(tem_current).astype(np.float)
        start_year = int(alllines[0][0])
        end_year = int(alllines[len(alllines)-1][0])
        latitude = float(alllines[0][2])
        distribution = indices.Distribution.gamma
        spei = indices.spei(scale=12, distribution=distribution, periodicity='monthly', data_start_year=start_year,
                            calibration_year_initial=start_year, calibration_year_final=end_year, precips_mm=precips_np,
                            temps_celsius=tem_np, latitude_degrees=latitude, pet_mm=None)
        for t in range(0, len(spei)):
            """ 1.假设数据没有遗漏，全年每个月都有 """
            # if t % 12 == 0:
            #     if t != 0:
            #         print("\n")
            #         fw.write('\n')
            #     print(alllines[t][0])
            #     fw.write(str(alllines[t][0]) + ' ')
            # print(spei[t])
            # fw.write(str(spei[t]) + ' ')

            """ 2.用英文逗号分割方便csv读取，带年.月好作图 """
            print(str(alllines[t][0]) + '年' + str(alllines[t][1]) + '月' + ' ' + str(spei[t]) + '\n')
            if str(spei[t]) != 'nan':
                fw.write(str(alllines[t][0]) + '-' + str(alllines[t][1]) + ',' + str(alllines[t][3]) + ','+
                         str(alllines[t][4]) + ',' + str(alllines[t][5]) + ',' + str(alllines[t][6]) + ','+
                         str(alllines[t][7]) + ',' +str(alllines[t][8]) + ',' + str(alllines[t][9]) + ','+ str(spei[t]) + '\n')
            else:
                continue
